#include <fstream>
#include <sstream>
#include <algorithm>
#include <numeric>
#include <stdexcept>
#include <vector>
#include <deque>
#include <string>
#include <iostream>
#include <istream>
#include <iomanip>

//#include <unistd.h>
//#include <sys/types.h>
//#include <sys/stat.h>
#include <boost/algorithm/string/replace.hpp>
#include <boost/foreach.hpp>

#define DEBUG_LEVEL 6
#include <debug_func.h>

#include <portable_timer.h>
#include <progress_indicator.h>
#include <index_string.h>
#define _MY_STL_MACROS
#include <stlhelper.h>
#include <ultoa.h>
#include <print_error.h>
#include <seq_pairs.h>
#include <open_fstream.h>
#include <get_filename.h>
#include <crc.h>
#include <digest.h>
#include "get_arg.h"

using boost::algorithm::replace_all;
using std::vector;
using std::deque;
using std::ofstream;
using std::ifstream;
using std::istream;
using std::istringstream;
using std::ostream;
using std::string;
using std::runtime_error;
using std::cerr;
using std::cin;
using std::cout;



//________________________________________________________________________________________

//		note_id

//________________________________________________________________________________________
void note_id(	const string&		id,
				t_index_string&		prot_ids,
				deque<unsigned>&	prot_to_gene_ids,
				deque<unsigned>&	cnt_prot_ids,
				unsigned&			prot_id)
{
	bool added;
	prot_id = prot_ids.index(id, added);

	//
	//	This protein no mapping to gene_id
	//
	if (added)
	{
		assert(index == cnt_prot_ids.size());
		assert(index == prot_to_gene_ids.size());
		// map to gene_id of unknown
		prot_to_gene_ids.push_back(0);
		cnt_prot_ids.push_back(0);
	}

}


//________________________________________________________________________________________

//		do_parse_pairs

//________________________________________________________________________________________
struct do_parse_pairs
{
	t_index_string&     		prot_ids;        
	deque<unsigned>&    		prot_to_gene_ids;
	deque<unsigned>&        	cnt_prot_ids;   
	deque<unsigned>&        	cnt_gene_ids;   
	bool						discard_self_self;
	std::string					comment_letters;
	unsigned					cnt_pairs;
	do_parse_pairs(	t_index_string&     			prot_ids_,        
					deque<unsigned>&    			prot_to_gene_ids_,
					deque<unsigned>&        		cnt_prot_ids_,   
					deque<unsigned>&        		cnt_gene_ids_,  
					bool							keep_self_self_, 
					const std::string&				comment_letters_):
			prot_ids(prot_ids_),         
			prot_to_gene_ids(prot_to_gene_ids_), 
			cnt_prot_ids(cnt_prot_ids_),    
			cnt_gene_ids(cnt_gene_ids_),    
			discard_self_self(!keep_self_self_),
			comment_letters(comment_letters_),
			cnt_pairs(0){}
	void operator()(istream& is, unsigned cnt_line)
	{
		// get data for one protein pair
		string name1, name2;

		// comment if false
		string comment_line;
		if (!seq_pairs::parse_ids(is, name1, name2, comment_line, comment_letters,
							  seq_pairs::ePARSE_WHOLE_LINE))
			return;
		++cnt_pairs;
		unsigned prot_id1;
		unsigned prot_id2;
		note_id(name1, prot_ids, prot_to_gene_ids, cnt_prot_ids, prot_id1);
		note_id(name2, prot_ids, prot_to_gene_ids, cnt_prot_ids, prot_id2);
		

		//
		// for counting protein and gene identifiers
		//
		if (prot_id1 == prot_id2 && discard_self_self)
			return;
		++cnt_prot_ids[prot_id1];
		++cnt_prot_ids[prot_id2];

		unsigned gene_id1 = prot_to_gene_ids[prot_id1];
		unsigned gene_id2 = prot_to_gene_ids[prot_id2];
		assert(gene_id1 < cnt_gene_ids.size());
		assert(gene_id2 < cnt_gene_ids.size());
		if (gene_id1 == gene_id2 && discard_self_self)
			return;
		++cnt_gene_ids[gene_id1];
		++cnt_gene_ids[gene_id2];
	}
};


		




//________________________________________________________________________________________

//		print_identifier_counts

//________________________________________________________________________________________
void print_identifier_counts(	ostream&						summary,
								ostream&						verbose,
								const t_index_string&			species,
								const string&					name,
								const vector<deque<unsigned> >&	identifiers_per_species)
{
	unsigned cnt_identifiers = 0;
	BOOST_FOREACH(const deque<unsigned>& ids, identifiers_per_species) 
		cnt_identifiers += ids.size();
	if (!cnt_identifiers)
		return;

	verbose << "\t" << 	std::setw(10)
		 << cnt_identifiers << "\tused " + name + " in total\n";
	summary
		 << cnt_identifiers << "\tused " + name + " in total\n";
	if (identifiers_per_species.size() > 1)
	{
		for (unsigned i = 0; i < identifiers_per_species.size(); ++i)
		{
			if (!identifiers_per_species[i].size())
				continue;
			verbose << "\t" << 	std::setw(10)
					<< identifiers_per_species[i].size() << "\tused " + name + " from species \t"
					<< species.string_at(i) << "\n";
			summary
					<< identifiers_per_species[i].size() << "\tused " + name + " from species \t"
					<< species.string_at(i) << "\n";
		}
	}
}

//________________________________________________________________________________________

//		count_ids_from_species

//________________________________________________________________________________________
void count_ids_from_species(const deque<unsigned>&		cnt_prot_ids,
							const deque<unsigned>&		cnt_gene_ids,
							const deque<unsigned>&		prot_to_gene_ids,
							const deque<unsigned>&		gene_ids_to_species,
							vector<deque<unsigned> >&	genes_per_species,
							vector<deque<unsigned> >&	prots_per_species)
{
	// count proteins
	for (unsigned i = 0; i < cnt_prot_ids.size(); ++i)
	{
		if (!cnt_prot_ids[i])
			continue;
		// convert prot to gene id to species id and increment
		unsigned species_id = gene_ids_to_species[prot_to_gene_ids[i]]; 
		prots_per_species[species_id].push_back(i);
	}

	// save genes
	for (unsigned i = 0; i < cnt_gene_ids.size(); ++i)
	{
		if (!cnt_gene_ids[i])
			continue;
		// convert gene id to species id and save
		genes_per_species[gene_ids_to_species[i]].push_back(i);
	}
}

//________________________________________________________________________________________

//		get_id_map

//________________________________________________________________________________________
void get_id_map(	istream& 			f_id_map,
					t_index_string&		prot_ids,
					t_index_string&		gene_ids,
					t_index_string&		species,
					deque<unsigned>&	prot_to_gene_ids,
					deque<unsigned>&	gene_ids_to_species,
					t_program_args&		args)
{
	unsigned cnt_line = 1;
	t_progress_indicator dots(args.VERBOSE(2),  10000);
	string prot, gene, species_str;


	// read next line
	f_id_map >> prot >> gene >> species_str;
	// if id is gene_id, does not need to be translated
	if (args.genes)
		prot = gene;

	
	while (f_id_map)
	{
		++dots;

		// index species
		unsigned species_id		= species.index(species_str);

		// index gene
		bool gene_added;
		unsigned gene_id		= gene_ids.index(gene, gene_added);
		if (gene_added)
			gene_ids_to_species.push_back(0);
		gene_ids_to_species[gene_id] = species_id;


		// index protein
		bool prot_added;
		unsigned prot_id		= prot_ids.index(prot, prot_added);
		if (prot_added)
			prot_to_gene_ids.push_back(0);
		prot_to_gene_ids[prot_id] = gene_id;

		// read next line
		f_id_map >> prot >> gene >> species_str;
		// if id is gene_id, does not need to be translated
		if (args.genes)
			prot = gene;
		++cnt_line;
	}

	dots.finish();
	args.VERBOSE(2) << "\t" << 	std::setw(10)  << dots.value() << "\tsequence identifiers read.\n";
	args.VERBOSE(2)	<< "\t" << 	std::setw(10)  << prot_ids.size()
					<< "\tused sequence identifiers read.\n";
	args.VERBOSE(2)	<< "\t" << 	std::setw(10)  << (gene_ids.size() - 1)
					<< "\tused gene identifiers read.\n";
	args.VERBOSE(2)	<< "\t" << 	std::setw(10)  << (species.size() - 1)
					<< "\tspecies.\n";
}

//8888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

//		Main

//8888888888888888888888888888888888888888888888888888888888888888888888888888888888888888
int main(int argc, char* argv[])
{
	t_program_args args;

	try
	{

	debug_func_cerr(5);

	// don't use C streams
	std::ios::sync_with_stdio(false);


	// get command line arguments
	if (!process_command_line_options(argc, argv, args))
		return 1;


	t_index_string		prot_ids;
	t_index_string		gene_ids;
	t_index_string		species;
	deque<unsigned>		prot_to_gene_ids;
	deque<unsigned>		gene_ids_to_species;


	//
	//	Default to unknown gene of unknown species
	// 		at index 0
	gene_ids.index("Unknown");
	gene_ids_to_species.push_back(0);
	species.index("Unknown");


	//
	//	Get Identifier maps
	//
	if (args.file_ids.length())
	{
		args.VERBOSE(2)	<< "\tObtaining sequence identifier map from ["
						<< get_lastdir_filename(args.file_ids) << "]...\n";

		ifstream f_id_map(args.file_ids.c_str());
		if (!f_id_map.good())
			throw std::runtime_error("Could not open " + args.file_ids);
		get_id_map(f_id_map, prot_ids, gene_ids, species, prot_to_gene_ids, gene_ids_to_species, args);
	}



	//
	//	Read through results and save
	//
	deque<unsigned>	cnt_prot_ids(prot_ids.size(), 0);
	deque<unsigned>	cnt_gene_ids(gene_ids.size(), 0);

	//
	// read sequence pairs
	//
	// filter alignment data
	//
	args.VERBOSE() << "\tReading sequence pair data from [" + 
						(args.file_input.length() ? 
						 get_lastdir_filename(args.file_input) : 
						 "STDIN") + "]\n";
	do_parse_pairs parse_pairs(	prot_ids,
								prot_to_gene_ids,
								cnt_prot_ids,
								cnt_gene_ids,
								args.keep_self_self,
								args.comment_letters);
	seq_pairs::read_data(parse_pairs, args.file_input, args.VERBOSE());

	args.VERBOSE() << "\t" << 	std::setw(10)
		 << parse_pairs.cnt_pairs << "\tsequence pairs in total\n";
	args.SUMMARY() << 	std::setw(10)
		 << parse_pairs.cnt_pairs << "\tsequence pairs in total\n";

	//
	//	Count the number of proteins per species
	//		and divide up genes into different species
	//
	vector<unsigned> 		 count_prots_per_species(species.size(), 0);
	vector<deque<unsigned> > genes_per_species(species.size());
	vector<deque<unsigned> > prots_per_species(species.size());
	count_ids_from_species(	cnt_prot_ids, 
							cnt_gene_ids,
							prot_to_gene_ids,
							gene_ids_to_species,
							genes_per_species,
							prots_per_species);


	//
	//	Print counts per species
	//		
	//
	if (args.genes)
	{
		print_identifier_counts(args.SUMMARY(), args.VERBOSE(), species, "genes", prots_per_species);
	}
	else
	{
		print_identifier_counts(args.SUMMARY(), args.VERBOSE(), species, "identifiers", prots_per_species);
		if (args.file_ids.length())
			print_identifier_counts(args.SUMMARY(), args.VERBOSE(), species, "genes", genes_per_species);
	}

	//
	//	Print gene identifiers if map
	//		
	//
	if (!args.genes						&&
		!args.path_output_gene.empty() &&
		gene_ids_to_species.size() > 1)
	{
		std::ofstream ostrm;
		std::string path;
		for (unsigned i = 0; i < species.size(); ++i)
		{
			if (!genes_per_species[i].size())
				continue;
			string new_path = args.path_output_gene;
			replace_all(new_path, "[SPECIES]", species.string_at(i));
			if (new_path != path)
			{
				ostrm.close();
				path = new_path;
				ostrm.open(path.c_str());
			}
			if (!ostrm)
				throw std::runtime_error("Could not open " + path);

			//
			//	just identifiers and the count
			//
			if (args.count_identifiers)
				BOOST_FOREACH(unsigned gene_id, genes_per_species[i]) 
					ostrm	<< gene_ids.string_at(gene_id)	<< "\t"
							<< cnt_gene_ids[gene_id]		<< "\n";

			//
			//	just print identifiers
			//
			else
				BOOST_FOREACH(unsigned gene_id, genes_per_species[i]) 
					ostrm	<< gene_ids.string_at(gene_id)	<< "\n";
		}
	}

	//
	//	Print protein identifiers
	//		
	//
	if (args.genes)
		args.path_output_prot = args.path_output_gene;
	if (!args.path_output_prot.empty())
	{
		// no map file
		if (gene_ids_to_species.size() <= 1)
		{
			string path = args.path_output_prot;
			replace_all(path, "[SPECIES]", "");
			ofstream ostrm(path.c_str());
			if (!ostrm)
				throw std::runtime_error("Could not open " + path);
			//
			//	just identifiers and the count
			//
			if (args.count_identifiers)
				BOOST_FOREACH(unsigned prot_id, prots_per_species[0]) 
					ostrm	<< prot_ids.string_at(prot_id)	<< "\t"
							<< cnt_prot_ids[prot_id]		<< "\n";

			//
			//	just print identifiers
			//
			else
				BOOST_FOREACH(unsigned prot_id, prots_per_species[0]) 
					ostrm	<< prot_ids.string_at(prot_id)	<< "\n";
		}
		else
		{
			std::ofstream ostrm;
			std::string path;
			for (unsigned i = 0; i < species.size(); ++i)
			{
				if (!prots_per_species[i].size())
					continue;
				string new_path = args.path_output_prot;
				replace_all(new_path, "[SPECIES]", species.string_at(i));
				if (new_path != path)
				{
					ostrm.close();
					path = new_path;
					ostrm.open(path.c_str());
				}
				if (!ostrm)
					throw std::runtime_error("Could not open " + path);

				//
				//	just identifiers and the count
				//
				if (args.count_identifiers)
					BOOST_FOREACH(unsigned prot_id, prots_per_species[i]) 
						ostrm	<< prot_ids.string_at(prot_id)	<< "\t"
								<< cnt_prot_ids[prot_id]		<< "\n";

				//
				//	just print identifiers
				//
				else
					BOOST_FOREACH(unsigned prot_id, prots_per_species[i]) 
						ostrm	<< prot_ids.string_at(prot_id)	<< "\n";
			}
		}
	}




	args.VERBOSE() << "\tAll finished.\n";

	}
	catch (std::exception& e)
	{
		std_print_error(args.CERR(), e.what());
		return 1;
	}

	return 0;
}

